#pragma once

#include <torch/torch.h>
#include <string>
#include <LibDL/utils.h>
#include <LibDL/options/dtype.h>
#include <LibDL/options/device.h>
#include <LibDL/ordered_dict/ordered_dict.h>
#include <LibDL/Tensor/Tensor.h>
#include <boost/any.hpp>
#include <boost/stacktrace.hpp>

namespace nn {
    class Module {
    protected:
        std::string _version = "1.3.0";

        std::shared_ptr<torch::nn::Module> core_;

        template<class T>
        std::shared_ptr<T> get_ptr() noexcept(false) {
            return std::dynamic_pointer_cast<T>(core_);
        }

    public:
        std::string version() const {
            return _version;
        }

        std::shared_ptr<torch::nn::Module> get_core() noexcept(false) {
            return core_;
        }

        std::shared_ptr<torch::nn::Module> get_core() const noexcept(false) {
            return core_;
        }

        static void print_module_pos(std::shared_ptr<torch::nn::Module> m) noexcept(false) {
            std::cout << m.get() << std::endl;
        }

        Module() noexcept(false) : core_(std::make_shared<torch::nn::Module>()) {
        }

        std::string to_string() const noexcept(false) {
            std::stringstream ss;
            core_->pretty_print(ss);
            return ss.str();
        }

        void register_module_(const std::string &name, std::shared_ptr<Module> module_holder)
        noexcept(false) {
            core_->register_module(name, module_holder->get_core());
        }

        Tensor register_parameter(const std::string &name, Tensor t, bool requires_grad = true)
        noexcept(false) {
            return core_->register_parameter(name, t, requires_grad);
        }

        Tensor register_buffer(const std::string &name, Tensor t)
        noexcept(false) {
            return core_->register_buffer(name, t);
        }


        void train() noexcept(false) {
            core_->train();
        }

        void eval() noexcept(false) {
            core_->eval();
        }

        std::vector<Tensor> parameters() const noexcept(false) {
            //return parameters();
            return utils::map_to_wrapper<Tensor>(core_->parameters());
        }

        std::vector<std::pair<std::string, std::shared_ptr<torch::nn::Module>>>
        named_modules() noexcept(false) {
            return core_->named_modules("", false).pairs();
        }

        std::vector<std::pair<std::string, nn::Module>> named_children() noexcept(false) {
            return utils::map_to_wrapper<std::pair<std::string, nn::Module>>(core_->named_children().pairs());
        }

        Module &to(const Device &d) noexcept(false) {
            core_->to(d);
            return *this;
        }

#if defined(SWIG)

        %javamethodmodifiers nn::Module::register_module_     "private";
        %proxycode
            %{
    public <T> T register_module(String s,T val) {
        if (val instanceof Module) {
            this.register_module_(s, (Module)val);
            return val;
        }
        System.out.println("Registering a non-module");
        return null;

    }

    public java.lang.String toString() {
        return this.to_string();
    }


    protected java.util.List<? extends LibDL.core.nn.Module> register_modules(String name,
            java.util.List<? extends LibDL.core.nn.Module> modules) {
        int i=0;
        for (LibDL.core.nn.Module model :modules){
            register_module(name+Integer.toString(i), model);
            i++;
        }
        return modules;
    }


    protected void register_modules(String[] names, java.util.List<? extends LibDL.core.nn.Module>... modules) {
        assert (names.length==modules.length);
        for (int i=0;i<names.length;i++){
            register_modules(names[i],modules[i]);
        }
    }

    protected void register_modules(String[] names, LibDL.core.nn.Module... modules) {
        assert (names.length==modules.length);
        for (int i=0;i<names.length;i++){
            register_module(names[i],modules[i]);
        }
    }

            %}
#endif //  SWIG

        explicit Module(std::shared_ptr<torch::nn::Module> core) noexcept(false)
                : core_(std::move(core)) {
        }
    };
}
